import numpy as np
import akshare as ak
import pandas as pd
from sqlalchemy import create_engine, text, exc
from datetime import datetime, timedelta
import time
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('stock_data.log'),
        logging.StreamHandler()
    ]
)

# PostgreSQL数据库配置
POSTGRES_SYSTEM_URL = "postgresql://leizhen:751982leizhen@localhost:5432/postgres"
DATABASE_URL = "postgresql://leizhen:751982leizhen@localhost/stock_db"
engine = create_engine(DATABASE_URL, pool_size=20, max_overflow=10)

def init_database():
    """初始化股票数据库表结构"""
    # 1. 检查数据库是否存在
    temp_engine = create_engine(POSTGRES_SYSTEM_URL)
    try:
        # 尝试连接到目标数据库
        test_engine = create_engine(DATABASE_URL)
        with test_engine.connect():
            pass
        logging.info("数据库已存在，继续初始化")
    except exc.OperationalError:
        logging.info("数据库不存在，开始创建...")
        try:
            # 创建新数据库
            with temp_engine.connect() as conn:
                conn.execute(text("COMMIT"))  # 结束可能存在的未完成事务
                conn.execute(text("CREATE DATABASE stock_db"))
                logging.info("数据库创建成功")
        except Exception as create_error:
            logging.error(f"创建数据库失败: {str(create_error)}")
            raise SystemExit("数据库创建失败，请检查：\n1. 用户权限\n2. 磁盘空间\n3. 网络连接")
    except exc.OperationalError as e:
        logging.error(f"数据库连接失败: {str(e)}")
        raise SystemExit("请检查：\n1. PostgreSQL服务是否运行\n2. 网络连接是否正常\n3. 用户名密码是否正确")
    except exc.ProgrammingError as e:
        logging.info("数据库已存在，继续初始化表结构")

    # 创建表结构
    with engine.connect() as conn:
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS stocks (
                symbol VARCHAR(10) PRIMARY KEY,
                name VARCHAR(100) NOT NULL,
                industry VARCHAR(50),
                listing_date DATE NOT NULL,
                last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """))
        
        # 创建基准表作为模板
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS stock_daily (
                id SERIAL PRIMARY KEY,
                symbol VARCHAR(10),
                trade_date DATE,
                open NUMERIC(10,2),
                close NUMERIC(10,2),
                high NUMERIC(10,2),
                low NUMERIC(10,2),
                volume BIGINT,
                amount NUMERIC(15,2),
                price_amplitude NUMERIC(5,2),
                change_pct NUMERIC(5,2),
                change_amount NUMERIC(10,2),
                turnover_rate NUMERIC(5,2)
            )
        """))
        conn.commit()

def get_all_stocks():
    """获取全量A股列表"""
    try:
        df = ak.stock_info_a_code_name()
        return df[['code', 'name']].rename(columns={'code': 'symbol', 'name': 'name'})
    except Exception as e:
        logging.error(f"获取股票列表失败: {str(e)}")
        return pd.DataFrame()

def get_stock_history(symbol, start_date):
    """获取个股历史行情数据（带重试机制）"""
    max_retries = 3
    for attempt in range(max_retries):
        try:
            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date.strftime("%Y%m%d"),
                end_date=datetime.now().strftime("%Y%m%d"),
                adjust="")
            return df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '股票代码': 'symbol',
                '成交额': 'amount',
                '振幅': 'price_amplitude',
                '涨跌幅': 'change_pct',
                '涨跌额': 'change_amount',
                '换手率': 'turnover_rate'
            })[['symbol', 'trade_date', 'open', 'close', 'high', 'low', 'volume', 
               'amount', 'price_amplitude', 'change_pct', 'change_amount', 'turnover_rate']]
        except Exception as e:
            if attempt == max_retries - 1:
                logging.warning(f"获取{symbol}数据失败: {str(e)}")
                return pd.DataFrame()
            time.sleep(2 ** attempt)
    return pd.DataFrame()

def main():
    init_database()
    
    # 获取全量股票列表
    stocks_df = get_all_stocks()
    if stocks_df.empty:
        logging.error("未获取到股票列表")
        return

    # 存储股票基本信息
    stocks_df.to_sql('stocks', engine, if_exists='replace', index=False)
    
    # 计算获取天数
    start_date = datetime.now() - timedelta(days=20)
    
    # 获取历史数据
    success = 0
    total = len(stocks_df)
    for idx, row in stocks_df.iterrows():
        symbol = row['symbol']
        name = row['name']
        
        logging.info(f"正在处理 [{idx+1}/{total}] {symbol}-{name}")
        
        history = get_stock_history(symbol, start_date)
        if not history.empty:
            # 确保列名与数据库schema完全匹配
            history = history[['symbol', 'trade_date', 'open', 'close', 'high', 'low', 'volume', 
                              'amount', 'price_amplitude', 'change_pct', 'change_amount', 'turnover_rate']].copy()
            history.drop_duplicates(subset=['symbol', 'trade_date'], inplace=True)

            # Check for existing records to avoid duplicates
            # 转换日期格式并确保类型一致
            # 转换日期格式并确保类型一致
            try:
                history['trade_date'] = pd.to_datetime(history['trade_date'], format='%Y-%m-%d')
                years = history['trade_date'].dt.year.unique()
                tables_to_check = [f'stock_daily_{year}' for year in years]
            except Exception as e:
                logging.error(f"日期格式转换失败: {str(e)}")
                logging.warning("跳过该股票数据处理")
                continue
            
            existing_dfs = []
            for table in tables_to_check:
                try:
                    # 仅查询与新数据年份对应的表
                    df = pd.read_sql(f"SELECT symbol, trade_date FROM {table}", engine)
                    existing_dfs.append(df)
                    logging.info(f"成功读取{len(df)}条{table}的现有记录")
                except Exception as e:
                    logging.warning(f"跳过{table}表查询: {str(e)}")
            
            existing = pd.concat(existing_dfs, ignore_index=True) if existing_dfs else pd.DataFrame(columns=['symbol', 'trade_date'])
            
            # 统一日期格式处理
            try:
                existing['trade_date'] = pd.to_datetime(existing['trade_date']).dt.date
                history['trade_date'] = pd.to_datetime(history['trade_date']).dt.date
            except Exception as e:
                logging.error(f"日期格式统一失败: {str(e)}")
                continue
            
            # 执行合并操作
            merged = history.merge(existing, on=['symbol', 'trade_date'], how='outer', indicator=True)
            new_records = merged[merged['_merge'] == 'left_only'].drop('_merge', axis=1)
            
            if not new_records.empty:
                # 使用SQLAlchemy核心接口进行批量插入
                with engine.connect() as db_conn:
                    # 数据清洗和类型转换
                    new_records = new_records.dropna()  # 移除空值记录
                    new_records['trade_date'] = pd.to_datetime(new_records['trade_date']).dt.date
                    
                    # 最终数据清洗方案（修复百分比和范围问题）
                    try:
                        # 1. 转换所有百分比字段
                        percent_cols = ['change_pct', 'turnover_rate', 'price_amplitude']
                        for col in percent_cols:
                            new_records[col] = pd.to_numeric(new_records[col], errors='coerce')
                            valid_mask = new_records[col].notna()
                            new_records.loc[valid_mask, col] = (
                                new_records.loc[valid_mask, col].astype(float) / 100.0
                            ).round(4)  # 保留4位小数避免舍入误差
                        
                        # 3. 数据范围过滤
                        new_records = new_records[
                            (new_records['change_pct'].between(-999.99, 999.99)) &
                            (new_records['turnover_rate'].between(0, 100)) &  # 换手率应≤100%
                            (new_records['price_amplitude'] >= 0) &
                            (new_records['volume'] > 0) &
                            (new_records['amount'] > 0)
                        ]
                        
                        # 4. 最终数据验证
                        if new_records.empty:
                            logging.warning("有效数据为空，跳过插入")
                            continue
                    except Exception as e:
                        logging.error(f"处理百分比值时发生错误: {str(e)}")
                        continue  # 跳过当前批次问题数据
                    
                    # 2. 数值范围限制（根据数据库字段精度）
                    numeric_limits = {
                        'price_amplitude': (-999.99, 999.99),
                        'change_pct': (-999.99, 999.99),
                        'turnover_rate': (0, 999.99),  # 换手率不应为负数
                        'change_amount': (-99999999.99, 99999999.99)
                    }
                    for col, (min_val, max_val) in numeric_limits.items():
                        new_records[col] = new_records[col].clip(lower=min_val, upper=max_val)
                    
                    # 3. 格式化小数位数
                    new_records = new_records.round({
                        'open': 2,
                        'close': 2,
                        'high': 2,
                        'low': 2,
                        'amount': 2,
                        'price_amplitude': 2,
                        'change_pct': 2,
                        'change_amount': 2,
                        'turnover_rate': 2
                    })
                    
                    # 4. 安全处理无穷大和NaN值
                    try:
                        # 确保所有数值列都是float类型
                        numeric_cols = ['open', 'close', 'high', 'low', 'volume', 'amount',
                                       'price_amplitude', 'change_pct', 'change_amount', 'turnover_rate']
                        new_records[numeric_cols] = new_records[numeric_cols].apply(pd.to_numeric, errors='coerce', axis=1)
                        
                        # 替换无穷大值为NaN并删除
                        new_records = new_records.replace([np.inf, -np.inf], np.nan)
                        new_records = new_records.dropna()
                        logging.info(f"有效数据记录数: {len(new_records)}")
                    except Exception as e:
                        logging.error(f"清理无效值时发生错误: {str(e)}")
                        logging.info("跳过当前问题数据批次")
                        continue
                    
                    # 执行批量插入并处理异常
                    # 按年份分组处理数据
                    new_records['year'] = pd.to_datetime(new_records['trade_date']).dt.year
                    
                    for year, year_group in new_records.groupby('year'):
                        # 创建年度表（如果不存在）
                        db_conn.execute(text(f"""
                            CREATE TABLE IF NOT EXISTS stock_daily_{year} (
                                id SERIAL PRIMARY KEY,
                                symbol VARCHAR(10),
                                trade_date DATE,
                                open NUMERIC(10,2),
                                close NUMERIC(10,2),
                                high NUMERIC(10,2),
                                low NUMERIC(10,2),
                                volume BIGINT,
                                amount NUMERIC(15,2),
                                price_amplitude NUMERIC(5,2),
                                change_pct NUMERIC(5,2),
                                change_amount NUMERIC(10,2),
                                turnover_rate NUMERIC(5,2),
                                UNIQUE(symbol, trade_date)
                            )
                        """))
                        
                        # 插入年度数据
                        db_conn.execute(
                            text(f"""
                                INSERT INTO stock_daily_{year} 
                                (symbol, trade_date, open, close, high, low, volume, 
                                 amount, price_amplitude, change_pct, change_amount, turnover_rate)
                                VALUES 
                                (:symbol, :trade_date, :open, :close, :high, :low, :volume,
                                 :amount, :price_amplitude, :change_pct, :change_amount, :turnover_rate)
                                ON CONFLICT (symbol, trade_date) DO NOTHING
                            """),
                            year_group.to_dict(orient='records')
                        )
                        
                        # 统计并记录年度表数据量
                        count_result = db_conn.execute(
                            text(f"SELECT COUNT(*) FROM stock_daily_{year}")
                        ).scalar()
                        logging.info(f"年度表 stock_daily_{year} 当前总记录数: {count_result}")
                    db_conn.commit()
                    logging.info(f"成功插入{len(new_records)}条新记录")
            else:
                logging.info(f"No new records for {symbol}")
            success += 1
            time.sleep(1)  # 控制请求频率

    logging.info(f"数据采集完成，共处理{success}/{total}只股票")

if __name__ == "__main__":
    main()
